﻿#pragma once

#include <cstddef>
#include <cstdio>
#include <format>
#include <intrin.h>
#include <map>
#include <new>
#include <string>
#include <vector>

namespace xy
{

namespace utils
{

struct memspy
{
    struct info
    {
        void *ptr;
        std::size_t size;
        std::string caller;
    };

#ifdef MEMSPY_IMPLEMENTATION
    inline memspy();
#else
    memspy() = default;
#endif // MEMSPY_IMPLEMENTATION

#ifdef MEMSPY_IMPLEMENTATION
    inline ~memspy();
#else
    ~memspy() = default;
#endif // MEMSPY_IMPLEMENTATION

    inline static bool enable = false;
    inline static bool enableInfo = false;
    inline static bool ignoreSTL = true;
    inline static int stackDepth = 6;
    inline static std::map<void *, info> allocated{};

#ifdef MEMSPY_IMPLEMENTATION
    inline static std::string stack_trace();
#else
    static std::string stack_trace()
    {
        return "";
    }
#endif // MEMSPY_IMPLEMENTATION
};

struct enable_gaurd
{
    bool wasEnabled;

    enable_gaurd()
    {
        wasEnabled = xy::utils::memspy::enable;
        xy::utils::memspy::enable = false;
    }

    ~enable_gaurd()
    {
        xy::utils::memspy::enable = wasEnabled;
    }

    operator bool() const
    {
        return wasEnabled;
    }
};

} // namespace utils

} // namespace xy

#ifdef MEMSPY_IMPLEMENTATION

#include <windows.h>

// 取消一些 windows 定义的宏
#undef near
#undef far
#undef min
#undef max
// 注意 windows 在最前面

#include <dbghelp.h>

#pragma comment(lib, "dbghelp.lib")

xy::utils::memspy::memspy()
{
    enable = true;

    // 初始化符号句柄
    SymSetOptions(SYMOPT_UNDNAME | SYMOPT_DEFERRED_LOADS);
    if (!SymInitialize(GetCurrentProcess(), NULL, TRUE))
        std::fprintf(stderr, "SymInitialize failed: %lu\n", GetLastError());
}

xy::utils::memspy::~memspy()
{
    enable = false;

    // 需要注意析构顺序，确保 memspy 在 allocated 之前析构，否则这里遍历 allocated 会导致未定义行为
    for (auto &&[ptr, info] : allocated)
        std::printf("Warning: memory leak at %p (caller: %s)\n", ptr, info.caller.c_str());

    SymCleanup(GetCurrentProcess());
}

std::string xy::utils::memspy::stack_trace()
{
    void *stack[100];
    USHORT frames = CaptureStackBackTrace(0, stackDepth, stack, NULL);

    SYMBOL_INFO *symbol = (SYMBOL_INFO *)malloc(sizeof(SYMBOL_INFO) + 256 * sizeof(char));
    symbol->MaxNameLen = 255;
    symbol->SizeOfStruct = sizeof(SYMBOL_INFO);

    // 从 1 开始，跳过当前函数
    std::string result = "";
    for (USHORT i = 1; i < frames; i++)
    {
        std::string func = "";
        DWORD64 displacement = 0;
        if (SymFromAddr(GetCurrentProcess(), (DWORD64)stack[i], &displacement, symbol))
        {
            if (std::string(symbol->Name).starts_with("operator new") ||
                std::string(symbol->Name).starts_with("operator delete"))
                continue;

            // 跳过标准库函数
            if (ignoreSTL && std::string(symbol->Name).starts_with("std::"))
                return "";

            func = std::format(" {} <-", symbol->Name);
        }
        else
            func = std::format("0x{:X}", (DWORD64)stack[i]);

        result += func;
    }
    result += " ... ";
    std::free(symbol);
    return result;
}

inline void *operator new(std::size_t size)
{
    using spy = xy::utils::memspy;

    // RAII 即使异常也能保证 enable 状态恢复
    xy::utils::enable_gaurd gaurd;
    void *ptr = std::malloc(size);
    if (gaurd && ptr)
    {
        std::string caller = spy::stack_trace();
        if (caller.size() > 0)
        {
            if (spy::enableInfo)
                std::printf("Allocated %d bytes at %p (caller: %s)\n", (int)size, ptr, caller.c_str());
            spy::allocated[ptr] = spy::info{ptr, size, caller};
        }
    }

    if (!ptr)
        throw std::bad_alloc{};
    return ptr;
}

inline void operator delete(void *ptr) noexcept
{
    using spy = xy::utils::memspy;

    // RAII 即使异常也能保证 enable 状态恢复
    xy::utils::enable_gaurd gaurd;
    if (gaurd && ptr)
    {
        std::string caller = spy::stack_trace();
        if (caller.size() > 0)
        {
            auto it = spy::allocated.find(ptr);
            if (it == spy::allocated.end())
                std::printf("Warning: trying to deallocate unallocated memory at %p (caller: %s)\n", ptr,
                            caller.c_str());
            else
            {
                if (spy::enableInfo)
                    std::printf("Deallocated %d bytes at %p (caller: %s)\n", (int)it->second.size, ptr, caller.c_str());
                spy::allocated.erase(it);
            }
        }
    }
    std::free(ptr);
}

#endif // MEMSPY_IMPLEMENTATION